from .grammar import NonTerminal, Terminal, PRODUCTIONS
from collections import defaultdict


class LL1Analyzer:
    """LL(1)分析器，计算FIRST和FOLLOW集"""

    def __init__(self):
        self.first_sets = defaultdict(set)
        self.follow_sets = defaultdict(set)
        self.predictive_table = {}

    def compute_first_sets(self):
        """计算FIRST集"""
        changed = True
        while changed:
            changed = False
            for non_terminal, productions in PRODUCTIONS.items():
                for production in productions:
                    old_size = len(self.first_sets[non_terminal])
                    self._compute_first_for_production(non_terminal, production)
                    if len(self.first_sets[non_terminal]) > old_size:
                        changed = True

        return self.first_sets

    def _compute_first_for_production(self, non_terminal, production):
        """为单个产生式计算FIRST集"""
        if production[0] == Terminal.EPSILON:
            self.first_sets[non_terminal].add(Terminal.EPSILON)
            return

        for symbol in production:
            if isinstance(symbol, Terminal):
                self.first_sets[non_terminal].add(symbol)
                break
            elif isinstance(symbol, NonTerminal):
                # 添加非ε的FIRST集
                for first_symbol in self.first_sets[symbol]:
                    if first_symbol != Terminal.EPSILON:
                        self.first_sets[non_terminal].add(first_symbol)

                # 如果当前符号的FIRST集不包含ε，停止
                if Terminal.EPSILON not in self.first_sets[symbol]:
                    break
        else:
            # 所有符号都包含ε，则添加ε
            self.first_sets[non_terminal].add(Terminal.EPSILON)

    def compute_follow_sets(self):
        """计算FOLLOW集"""
        # 初始化，为开始符号添加$
        self.follow_sets[NonTerminal.PROGRAM].add(Terminal.EOF)

        changed = True
        while changed:
            changed = False
            for non_terminal, productions in PRODUCTIONS.items():
                for production in productions:
                    for i, symbol in enumerate(production):
                        if isinstance(symbol, NonTerminal):
                            old_size = len(self.follow_sets[symbol])

                            # 规则1: 如果A → αBβ，则FIRST(β) - {ε} ⊆ FOLLOW(B)
                            if i + 1 < len(production):
                                next_symbol = production[i + 1]
                                if isinstance(next_symbol, Terminal):
                                    self.follow_sets[symbol].add(next_symbol)
                                else:
                                    for first_sym in self.first_sets[next_symbol]:
                                        if first_sym != Terminal.EPSILON:
                                            self.follow_sets[symbol].add(first_sym)

                            # 规则2: 如果A → αB 或 A → αBβ且β →* ε，则FOLLOW(A) ⊆ FOLLOW(B)
                            if i + 1 == len(production) or self._can_derive_epsilon(production[i + 1:]):
                                for follow_sym in self.follow_sets[non_terminal]:
                                    self.follow_sets[symbol].add(follow_sym)

                            if len(self.follow_sets[symbol]) > old_size:
                                changed = True

        return self.follow_sets

    def _can_derive_epsilon(self, symbols):
        """判断符号序列是否能推导出ε"""
        for symbol in symbols:
            if isinstance(symbol, Terminal):
                return symbol == Terminal.EPSILON
            elif Terminal.EPSILON not in self.first_sets[symbol]:
                return False
        return True

    def build_predictive_table(self):
        """构建预测分析表"""
        for non_terminal, productions in PRODUCTIONS.items():
            for production in productions:
                first_alpha = self._get_first_of_sequence(production)

                for terminal in first_alpha:
                    if terminal != Terminal.EPSILON:
                        key = (non_terminal, terminal)
                        if key in self.predictive_table:
                            raise Exception(f"Grammar is not LL(1): conflict at {key}")
                        self.predictive_table[key] = production

                if Terminal.EPSILON in first_alpha:
                    for terminal in self.follow_sets[non_terminal]:
                        key = (non_terminal, terminal)
                        if key in self.predictive_table:
                            raise Exception(f"Grammar is not LL(1): conflict at {key}")
                        self.predictive_table[key] = production

        return self.predictive_table

    def _get_first_of_sequence(self, sequence):
        """获取符号序列的FIRST集"""
        first_set = set()

        for symbol in sequence:
            if isinstance(symbol, Terminal):
                first_set.add(symbol)
                break
            else:
                first_set.update(self.first_sets[symbol] - {Terminal.EPSILON})
                if Terminal.EPSILON not in self.first_sets[symbol]:
                    break
        else:
            first_set.add(Terminal.EPSILON)

        return first_set